import dashscope
from dashscope import TextEmbedding
from dashvector import Client

# Configure the API keys and client setup
dashscope.api_key = 'sk-39e47712f94a4816a18f356dd26aa1c8'
client = Client(
    api_key='sk-hz4xtal1OMo24Ft7hLC1EQdSZimEw6918B4CE47C911EF90AA2A0812EF32E2',
    endpoint='vrs-cn-j4g3ukace0004v.dashvector.cn-hangzhou.aliyuncs.com'
)


def generate_embeddings(text):
    rsp = TextEmbedding.call(model=TextEmbedding.Models.text_embedding_v1, input=text)
    embeddings = [record['embedding'] for record in rsp.output['embeddings']]
    return embeddings if isinstance(text, list) else embeddings[0]


def query_collection(collection_type, questions, endpoint):
    # Configure the API keys and client setup
    dashscope.api_key = 'sk-39e47712f94a4816a18f356dd26aa1c8'
    client = Client(
        api_key='sk-hz4xtal1OMo24Ft7hLC1EQdSZimEw6918B4CE47C911EF90AA2A0812EF32E2',
        endpoint=endpoint
    )
    results = []

    if collection_type == 'Picture_library':
        collection = client.get('Picture_library')
        output_fields = ['text_content', 'Material_url', 'Material_id']
    elif collection_type == 'Video_library':
        collection = client.get('Video_library')
        output_fields = ['text_content', 'Material_url', 'Video_length', 'Material_id']
    elif collection_type == 'Audio_library':
        collection = client.get('Audio_library')
        output_fields = ['text_content', 'Material_url', 'audio_length', 'Material_id']
    else:
        raise ValueError(f"Unsupported collection type: {collection_type}")

    for question in questions:
        embeddings = generate_embeddings(question)
        rsp = collection.query(embeddings, output_fields=output_fields, topk=1)

        for doc in rsp.output:
            result = {
                'id': doc.id,
                'text_content': doc.fields['text_content'],
                'Material_url': doc.fields['Material_url'],
                'Material_id': doc.fields['Material_id'],

            }
            if collection_type == 'Video_library':
                result['Video_length'] = doc.fields['Video_length']
            elif collection_type == 'Audio_library':
                result['audio_length'] = doc.fields['audio_length']
            elif collection_type == 'Picture_library':
                pass

            results.append(result)

    return results

# # Example usage
# questions = ["迪丽热巴","Sample text for audio"]
# collection_type = 'Audio_library'
#
# results = query_collection(collection_type, questions)
# print(results)
